from typing import ClassVar, override

import numpy as np
from pyecharts import options as opts
from pyecharts.charts import Scatter3D

from ...processors import ClusterStage
from ...utils import OutputManager
from ..interface import Marshaler

__all__ = ["HTMLClusterMarshaler"]


class HTMLClusterMarshaler(Marshaler[ClusterStage]):
    OUTPUT_NAME: ClassVar[str] = "kmeans_clusters.html"

    _output_manager: OutputManager
    _kmeans_config: dict
    _config: dict

    def __init__(self, output_manager: OutputManager, global_config: dict) -> None:
        self._output_manager = output_manager
        self._kmeans_config = global_config["clusterizer"]["kmeans"]
        self._config = global_config["clusterizer"]["kmeans"]["html"]

    @override
    def marshal(self, stage: ClusterStage) -> None:
        scatter3d = Scatter3D(
            init_opts=opts.InitOpts(
                width=self._config["width"],
                height=self._config["height"],
                theme=self._config["theme"],
            )
        )

        # Add data points for each cluster
        for i in range(self._kmeans_config["n_clusters"]):
            cluster_points = stage.points[stage.labels == i]

            # Add jitter effect to scatter points
            jittered_points = cluster_points.copy()
            if len(jittered_points) > 0:
                jittered_points += np.random.normal(0, self._kmeans_config["jitter_scale"], jittered_points.shape)
                # Ensure the jittered values are within the valid range
                jittered_points = np.clip(jittered_points, 0, 1)

            data = []
            for point in jittered_points:
                data.append([float(point[0]), float(point[1]), float(point[2])])

            scatter3d.add(
                series_name=f"Cluster {i + 1}",
                data=data,
                itemstyle_opts=opts.ItemStyleOpts(color=self._config["colors"][i % len(self._config["colors"])]),
            )

        # Add cluster centers
        centers_data = []
        for center in stage.centers:
            centers_data.append([float(center[0]), float(center[1]), float(center[2])])

        scatter3d.add(
            series_name="Centroids",
            data=centers_data,
            itemstyle_opts=opts.ItemStyleOpts(color=self._config["centroid_color"]),
        )

        # Configure the chart options
        scatter3d.set_global_opts(
            title_opts=opts.TitleOpts(title="User Clusters (3D)"),
            visualmap_opts=opts.VisualMapOpts(
                is_show=False,
                max_=1.0,
                min_=0.0,
            ),
            tooltip_opts=opts.TooltipOpts(
                formatter="""
                <div>
                    <p>{a}</p>
                    <p>Neutral: {c[0]}</p>
                    <p>Positive: {c[1]}</p>
                    <p>Negative: {c[2]}</p>
                </div>
                """
            ),
            legend_opts=opts.LegendOpts(is_show=True),
        )

        # Save the HTML file
        scatter3d.render(str(self._output_manager.filepath(HTMLClusterMarshaler.OUTPUT_NAME)))
